# Single-source shortest path in complete graphs. 

abstract type CompleteGraphShortestPathSolver end
function build!(solver::CompleteGraphShortestPathSolver, reward::Matrix{Distribution}, source::Int, destination::Int)
  nothing
end

struct CompleteGraphShortestPath <: CombinatorialInstance
  # Mandatory properties. 
  n_arms::Int

  # Probability distributions for the arm rewards. 
  reward::Matrix{Distribution}
  optimal_reward::Float64 # Cache. 

  # Characteristics of the path. 
  source::Int
  destination::Int

  # Internal solver. 
  solver::CompleteGraphShortestPathSolver

  function CompleteGraphShortestPath(reward::Matrix{Distribution}, source::Int, destination::Int, solver::CompleteGraphShortestPathSolver)
    n = size(reward, 1)
    if n != size(reward, 2)
      error("Weighted adjacency matrix of rewards must be square.")
    end

    if source <= 0
      error("Source node is not an index within the weighted adjacency matrix: it is negative or zero.")
    end
    if source > n
      error("Source node is not an index within the weighted adjacency matrix: it is too large.")
    end
    if destination <= 0
      error("Destination node is not an index within the weighted adjacency matrix: it is negative or zero.")
    end
    if destination > n
      error("Destination node is not an index within the weighted adjacency matrix: it is too large.")
    end

    # Prepare the solver to be used (if required). 
    build!(solver, reward, source, destination)

    # Precompute certain results. 
    weights = Dict((i, j) => mean(reward[i, j]) for i in 1:n, j in 1:n)
    path = solve_linear(solver, weights)
    optimal_reward = sum(weights[(edge[1], edge[2])] for edge in path)

    # Done! 
    return new(n * (n - 1) / 2, reward, optimal_reward, source, destination, solver)
  end
end

function initial_state(instance::CompleteGraphShortestPath) 
  n = size(instance.reward, 1)
  zeroCounts = Dict((i, j) => 0 for i in 1:n, j in 1:n)
  zeroRewards = Dict((i, j) => 0.0 for i in 1:n, j in 1:n)
  return State{Tuple{Int, Int}}(0, 0.0, 0.0, zeroCounts, zeroRewards, copy(zeroRewards))
end

function initial_trace(instance::CompleteGraphShortestPath) 
  return Trace{Tuple{Int, Int}}(State{Tuple{Int, Int}}[], Vector{Tuple{Int, Int}}[], Float64[], PolicyDetails[])
end

function pull(instance::CompleteGraphShortestPath, arms::Vector{Tuple{Int, Int}}) 
  return Float64[rand(instance.reward[arm[1], arm[2]]) for arm in arms]
end

solve_linear(instance::CompleteGraphShortestPath, weights::Dict{Tuple{Int, Int}, Float64}) = solve_linear(instance.solver, weights)

function is_feasible(instance::CompleteGraphShortestPath, arms::Vector{Tuple{Int, Int}}) 
  path_length = length(arms)

  if path_length > instance.n_arms
    return false
  end
  
  if path_length < 1
    return false
  end
  
  n = size(instance.reward, 1)

  # Check whether the path goes from the source to the destination. Think about the case where the path is just one edge. 
  for (i, arm) in enumerate(arms)
    # The first edge must start at the source. 
    if i == 1 && arm[1] != instance.source
      return false
    end
    
    # The last edge must end at the source. 
    if i == n && arm[2] != instance.destination
      return false
    end

    # If after the first edge, this edge must be adjacent to the same node as the previous one. 
    if i > 1 && arms[i - 1][2] != arm[1]
      return false
    end
  end

  return true
end
